package finaloutput

import (
	"context"
	"gitee.com/dn-jinmin/langchaingo"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/prompts"
)

type FinalOutput struct {
	chains.Chain
}

func NewFinalOutput(llm llms.Model) *FinalOutput {
	return &FinalOutput{
		Chain: chains.NewLLMChain(llm, prompts.NewPromptTemplate(_defaultFinalOutputPrompts, []string{langchaingo.Input})),
	}
}

func (f *FinalOutput) Call(ctx context.Context, input string) (string, error) {
	return chains.Predict(ctx, f.Chain, map[string]any{
		langchaingo.Input: input,
	})
}
